#include "faiss_ivfflat_index.h"
#include <faiss/index_io.h>
#include <filesystem>

FaissIVFFLATIndex::FaissIVFFLATIndex(const ConfigParser &config) {
    nprobe = config.get_int("nprobe", 145);
    std::string metric_str = config.get_string("metric", "L2");
    if (metric_str == "IP") {
        metric = faiss::METRIC_INNER_PRODUCT;
    } else {
        metric = faiss::METRIC_L2;
    }
    index_path = config.get_string("index_path", "index.faiss");
    save_or_load = config.get_string("save_or_load", "neither");
}

FaissIVFFLATIndex::~FaissIVFFLATIndex() {
    index.reset();
    quantizer.reset();
}

void FaissIVFFLATIndex::build(const Datasets *data) {
    if (save_or_load == "load") {
        auto loaded_index = faiss::read_index(index_path.c_str());
        auto ivfflat_ptr = dynamic_cast<faiss::IndexIVFFlat *>(loaded_index);
        if (ivfflat_ptr) {
            index.reset(ivfflat_ptr);
            index->own_fields = true;
            quantizer.reset();
        } else {
            delete loaded_index;
            throw std::runtime_error("Loaded index is not IVFFlat");
        }
    } else {
        int nlist = int(4 * sqrt(data->nb));
        if (metric == faiss::METRIC_L2) {
            quantizer = std::make_unique<faiss::IndexFlatL2>(data->d);
        } else {
            quantizer = std::make_unique<faiss::IndexFlatIP>(data->d);
        }
        index = std::make_unique<faiss::IndexIVFFlat>(quantizer.release(), data->d, nlist);
        index->train(data->nb, data->xb);
        index->add(data->nb, data->xb);
        if (save_or_load == "save") {
            std::filesystem::path dir_path = std::filesystem::path(index_path).parent_path();
            if (!dir_path.empty() && !std::filesystem::exists(dir_path)) {
                std::filesystem::create_directories(dir_path);
            }
            faiss::write_index(index.get(), index_path.c_str());
        }
    }
}

void FaissIVFFLATIndex::prepare_search() {
    ivf_params.nprobe = nprobe;
}

void FaissIVFFLATIndex::search_single(const float *xq, int k, float *distances, int64_t *labels) {
    int n = 1;
    index->search(n, xq, k, distances, labels, &ivf_params);
}

void FaissIVFFLATIndex::search_batch(const int nq, const float *xq, int k, float *distances, int64_t *labels) {
    index->search(nq, xq, k, distances, labels, &ivf_params);
}
